#! /usr/bin/env python
# Create contour plots of a spectrum with each peak in limits.in labeled

import nmrglue as ng
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm

# plot parameters
cmap = matplotlib.cm.Blues_r    # contour map (colors to use for contours)
contour_start   = 30000     # contour level start value
contour_num     = 20        # number of contour levels
contour_factor  = 1.20      # scaling factor between contour levels
textsize = 6       # text size of labels

# calculate contour levels
cl = [contour_start*contour_factor**x for x in xrange(contour_num)]

# read in the data from a NMRPipe file
dic,data = ng.pipe.read("../../common_data/2d_pipe/test.ft2")

# read in the integration limits
peak_list = np.recfromtxt("limits.in")

# create the figure
fig = plt.figure()
ax = fig.add_subplot(111)

# plot the contours
ax.contour(data,cl,cmap=cmap,extent=(0,data.shape[1]-1,0,data.shape[0]-1))

# loop over the peaks
for name,x0,y0,x1,y1 in peak_list:

    if x0>x1:
        x0,x1 = x1,x0
    if y0>y1:
        y0,y1 = y1,y0

    # plot a box around each peak and label
    ax.plot([x0,x1,x1,x0,x0],[y0,y0,y1,y1,y0],'k')
    ax.text(x1+1,y0,name,size=textsize,color='r')
    
# set limits
ax.set_xlim(1900,2200)
ax.set_ylim(750,1400)

# save the figure
fig.savefig("assignments.png")
